#!/usr/bin/env fish

# 该脚本使用 deepseek 自动生成, 通过两轮交互得到, 使用的 prompt 如下:
# * Round 1:
#   我有一个命令行工具(`f2h5`)，它接收多个输入文件和一个输出文件。该工具这多个输入文件的数据进行转换，将转换的结果输出到输出文件中。
#   现在有一个输入目录和一个输出目录，我需要将输入目录中的文件，以每 5 个为一组，输入给 `f2h5` 工具，
#   并输出到输出目录中。输出文件以序号命名，零前缀填充补齐为 4 个数字。`f2h5` 的调用例子如下：
#   ```
#   f2h5 --output_file 0001.h5 --fits_files=FRB_1.fits --fits_files+=FRB_2.fits --fits_files+=FRB_3.fits --fits_files+=FRB_4.fits --fits_files+=FRB_5.fits
#   ```
#   请使用 fish shell 语言，写一个脚本，完成该任务。
# * Round 2:
#   上面的脚本，输入文件分组的大小固定为 5。请将它修改为脚本的输入参数，该参数为可选的，如果用户没有设置，默认为 8.
# * Round 3:
#   现在，我有一个新的需求：当已经完成部分转换，需要继续进行转换时，目录中已经存在一些转换好的 HDF5 文件(.h5 后缀），需要先统计已经完成转换的部分，然后再继续进行转换。
#   每个 HDF5 文件中包含的 fits 文件列表信息可以使用 `h5ls` 命令给出。例如：
#
#     ```bash
#     $ h5ls 0001.h5 | cut -d' ' -f1
#     FRB20240114A_tracking-M01_0001
#     FRB20240114A_tracking-M01_0002
#     FRB20240114A_tracking-M01_0003
#     FRB20240114A_tracking-M01_0004
#     FRB20240114A_tracking-M01_0005
#     FRB20240114A_tracking-M01_0006
#     FRB20240114A_tracking-M01_0007
#     FRB20240114A_tracking-M01_0008
#     ```
#
#   注意以上列出的文件名不包含 .fits 后缀。注意，继续进行转换时，输出文件的编号应该继续编号，而非从头开始编号。

set input_dir $argv[1]
set output_dir $argv[2]
set group_size 8  # 默认分组大小

# 参数检查和处理
if test (count $argv) -lt 2 -o (count $argv) -gt 3
    echo "Usage: ./process_fits.fish input_dir output_dir [group_size=8]"
    exit 1
end

# 处理可选的分组大小参数
if test (count $argv) -eq 3
    set group_size $argv[3]
    if not string match -qr '^[1-9]\d*$' -- $group_size
        echo "Error: Invalid group size '$group_size'. Must be a positive integer"
        exit 1
    end
end

# 检查输入目录
if not test -d $input_dir
    echo "Error: Input directory does not exist"
    exit 1
end

# 创建输出目录
mkdir -p $output_dir

# 获取原始文件列表并排序
set fits_files (ls $input_dir/*.fits | sort -V)

# 步骤1：收集已存在的 HDF5 文件信息
set max_group_num 0
set processed_fits
set verified_h5files

# 查找最大组号并收集已处理的 FITS 文件
for h5file in $output_dir/*.h5
    # 提取已处理的 FITS 文件
    if h5ls $h5file >/dev/null 2>&1
        set -a verified_h5files $h5file
        set datasets (h5ls $h5file | cut -d' ' -f1)
        for dataset in $datasets
            set -a processed_fits "$dataset.fits"
        end
    else
        echo "Warning: 无法读取 HDF5 文件 $h5file，可能已损坏. 将进行删除，并重新生成"
        rm -f $h5file
    end
end

for h5file in $verified_h5files
    # 提取组号
    set base (basename $h5file .h5)
    if string match -qr '^\d{4}$' -- $base
        set group_num (math "$base + 0")
        test $group_num -gt $max_group_num && set max_group_num $group_num
    end
end

# 步骤2：过滤已处理的文件
set remaining_files
for file in $fits_files
    set base_file (basename $file)
    if not contains $base_file $processed_fits
        set -a remaining_files $file
    end
end

# 计算剩余文件的分组信息
set total_files (count $remaining_files)
set group_count (math -s0 "ceil($total_files / $group_size)")

# 步骤3：执行转换
for group_num in (seq 1 $group_count)
    # 计算实际组号
    set actual_group_num (math "$max_group_num + $group_num")

    # 计算文件范围
    set start_idx (math "$group_size * ($group_num - 1) + 1")
    set end_idx (math "$group_size * $group_num")

    # 处理最后一组
    if test $end_idx -gt $total_files
        set end_idx $total_files
    end

    # 构建参数列表
    set cmd_args
    for idx in (seq $start_idx $end_idx)
        set file $remaining_files[$idx]
        if test $idx -eq $start_idx
            set -a cmd_args --fits_files=$file
        else
            set -a cmd_args --fits_files+=$file
        end
    end

    # 生成输出文件名
    set outfile (printf "%s/%04d.h5" $output_dir $actual_group_num)

    # 执行转换命令
    set rest_files (math "$total_files - $end_idx")
    echo "Processing group $actual_group_num ($start_idx-$end_idx, 剩余 fits 文件数： $rest_files)"
    f2h5 --output_file $outfile $cmd_args
    or begin
        echo "Error occurred during processing group $actual_group_num with: f2h5 --output_file $outfile $cmd_args"
        echo "Removing $outfile"
        if test -f $outfile
            rm -f $outfile
        end
        exit 1
    end
end

# 最终统计
echo "Processing completed!"
echo "  原始文件总数: "(count $fits_files)
echo "  已跳过文件数: "(count $processed_fits)
echo "  本次处理文件数: "(count $remaining_files)
echo "  生成新组数量: $group_count"
